#include <stdlib.h>
#include <stdio.h>
#include <math.h>
#include <ctype.h>
#include <unistd.h>
#include <string.h>

#include "mpi.h"

int n;
int method;
int w;
int nWithBorders;
int mWithBorders;
int nodeedge; /* a task works on a nodeedge x nodeedge matrix */
int nblock; /* number of tasks per row of matrix            */
int nproc;/* total number of tasks (processors)           */
MPI_Comm cartcomm;
int coords[2];

int main(int argc, char **argv);
void doblack(double M[][nodeedge + 2]);
void dored(double M[][nodeedge + 2]);
void dojacobi(double M[][nodeedge + 2]);
void exchange(double M[][nodeedge + 2], int comm[], int rank);
void iterate(double M[][nodeedge + 2], double result[][n], int rank, int comm[]);
void timeIteration(double M[][nodeedge + 2], double result[][n], int rank,
		int comm[], int step, int finish, int coords[]);
void setcomm(int coords[], int dims[], int comm[]);
void setex(double ex[], double M[][nodeedge + 2], int which);
void unpack(double M[][nodeedge + 2], int where, double in[]);
float heatFunction(int x, int y, int t);

int main(int argc, char **argv) {
	int comm[4];
	FILE *fp;
	int i;
	int j;
	int ntasks;
	int rank, c;
	int step;
	int finish;
	double wtime;
	char* filename;

	MPI_Init(&argc, &argv);

	char hostname[256];
	gethostname(hostname, 256);

	while ((c = getopt(argc, argv, "f:n:m:w:s:e:")) != -1) {
		switch (c) {
		case 'f':
			filename = optarg;
			break;
		case 'n':
			nWithBorders = atoi(optarg);
			break;
		case 'm':
			//0 - Jacobi, 1 - RB , 2- SOR, 3 - Timed
			method = atoi(optarg);
			break;
		case 's':
			step = atoi(optarg);
			break;
		case 'e':
			finish = atoi(optarg);
			break;
		case 'w':
			w = atof(optarg);
			break;
		default:
			fprintf(stderr, "Bad args.\n", optopt);
			return -1;
		}
	}
	if (method != 2)
		w = 1.0;
	n = nWithBorders - 2;
	double result[nWithBorders][nWithBorders];
	wtime = MPI_Wtime();

	MPI_Comm_size(MPI_COMM_WORLD, &ntasks);
	MPI_Comm_rank(MPI_COMM_WORLD, &rank);

//	if(rank == 0){
	FILE *f = fopen(filename, "rt");
	for (i = 0; i < nWithBorders; i++) {
		for (j = 0; j < nWithBorders; j++) {
			fscanf(f, "%lf ", &result[i][j]);
		}
		fscanf(f, "\n");
	}

	fclose(f);
//	}
	if (rank == 0) {
		for (i = 0; i < nWithBorders; i++) {
			for (j = 0; j < nWithBorders; j++) {
				printf("%lf ", result[i][j]);
			}
			printf("\n");
		}
	}

	int ndims = 2;
	int dims[2] = { 0, 0 };
	MPI_Dims_create(ntasks, ndims, dims);
	int periods[2] = { 0, 0 };
	int reorder = 0;
	MPI_Cart_create(MPI_COMM_WORLD, ndims, dims, periods, reorder, &cartcomm);

	MPI_Comm_rank(cartcomm, &rank);
	MPI_Cart_coords(cartcomm, rank, ndims, coords);

//	int sizes[2] = {n,n};
//	int subsizes[2] = {n/procgridsize, n/procgridsize};
//
//	if(rank == 0){
//		MPI_Scatterv(global, counts, displs, /* proc i gets counts[i] types from displs[i] */
//		            resizedtype,
//		            local, 3*3, MPI_INT;   /* I'm receiving 3*3 MPI_INTs into local */
//		            0, MPI_COMM_WORLD);
//	}

// part of the main matrix for process with given rank
	int verticalN = (double) n / (double) dims[0];
	int verticalStart = coords[0] * verticalN;
	int verticalEnd = verticalStart + verticalN;

	int horizontalN = (double) n / (double) dims[1];
	int horizontalStart = coords[1] * horizontalN;
	int horizontalEnd = horizontalStart + horizontalN;

	//TODO: This supports only square matrixes
	nodeedge = verticalN;
	nblock = n / nodeedge;
	nproc = nblock * nblock;
	double M[verticalN + 2][horizontalN + 2];

	int k, m;
	for (i = verticalStart, k = 0; i <= verticalEnd + 1; i++, k++) {
		for (j = horizontalStart, m = 0; j <= horizontalEnd + 1; j++, m++) {
			M[k][m] = result[i][j];
		}
	}
	setcomm(coords, dims, comm);

	char out[512] = "";
	sprintf(out, "rank %d@%s on (%d,%d) comm: %d %d %d %d v:%d-%d, h:%d-%d\n",
			rank, hostname, coords[0], coords[1], comm[0], comm[1], comm[2],
			comm[3], verticalStart, verticalEnd, horizontalStart,
			horizontalEnd);

	for (k = 0; k <= verticalN + 1; k++) {
		for (m = 0; m <= horizontalN + 1; m++) {
			char doubleStr[10];
			sprintf(doubleStr, "%lf", M[k][m]);
			strcat(out, doubleStr);
			strcat(out, " ");
		}
		strcat(out, "\n");
	}

	printf("\n%s\n", out);

//	double M[][nodeedge + 2]; // TODO watch this
	if (method != 3) {
		iterate(M, result, rank, comm);

		wtime = MPI_Wtime() - wtime;

		printf("  Task %i took %6.3f seconds\n", rank, wtime);

		char out2[256] = "";
		if (rank == 0) {
			for (i = 0; i < nWithBorders; i++) {
				for (j = 0; j < nWithBorders; j++) {
					char doubleStr2[10];
					sprintf(doubleStr2, "%lf", result[i][j]);
					strcat(out2, doubleStr2);
					strcat(out2, " ");
				}
				strcat(out2, "\n");
			}
			printf("\n%s\n", out2);
		}
	} else {
timeIteration	(M, result, rank, comm, step, finish, coords );
}
	MPI_Finalize();
	return 0;
}

float heatFunction(int x, int y, int t) {
//	return (x == nWithBorders/2 && y == nWithBorders/2) ? 50.0 : 0.0;
	return y == 1 && x == t/5 ? 50.0 : 0.0;
}

void doblack(double M[][nodeedge + 2])

/******************************************************************************/
/*
 DOBLACK iterates on the upper right and lower left corners of my matrix.
 */
{
	int i;
	int j;
	/*
	 Upper right corner.
	 */
	for (i = 1; i <= nodeedge / 2; i++) {
		for (j = nodeedge / 2 + 1; j <= nodeedge; j++) {
			M[i][j] = w / 4.0
					* (M[i - 1][j] + M[i][j - 1] + M[i + 1][j] + M[i][j + 1])
					+ (1.0 - w) * M[i][j];
		}
	}
	/*
	 Lower left corner.
	 */
	for (i = nodeedge / 2 + 1; i <= nodeedge; i++) {
		for (j = 1; j <= nodeedge / 2; j++) {
			M[i][j] = w / 4.0
					* (M[i - 1][j] + M[i][j - 1] + M[i + 1][j] + M[i][j + 1])
					+ (1.0 - w) * M[i][j];
		}
	}
	return;
}

void dojacobi(double M[][nodeedge + 2]) {
	int i;
	int j;
	for (i = 1; i <= nodeedge; i++) {
		for (j = 1; j <= nodeedge; j++) {
			M[i][j] = (M[i - 1][j] + M[i][j - 1] + M[i + 1][j] + M[i][j + 1])
					/ 4.0;
		}
	}
}

void dored(double M[][nodeedge + 2])

/*
 DORED iterates on the upper left and lower right corners of my matrix.
 */
{
	int i;
	int j;
	/*
	 Upper left corner.
	 */
	for (i = 1; i <= nodeedge / 2; i++) {
		for (j = 1; j <= nodeedge / 2; j++) {
			M[i][j] = w / 4.0
					* (M[i - 1][j] + M[i][j - 1] + M[i + 1][j] + M[i][j + 1])
					+ (1.0 - w) * M[i][j];
		}
	}
	/*
	 Lower right corner.
	 */
	for (i = nodeedge / 2 + 1; i <= nodeedge; i++) {
		for (j = nodeedge / 2 + 1; j <= nodeedge; j++) {
			M[i][j] = w / 4.0
					* (M[i - 1][j] + M[i][j - 1] + M[i + 1][j] + M[i][j + 1])
					+ (1.0 - w) * M[i][j];
		}
	}
	return;
}

void exchange(double M[][nodeedge + 2], int comm[], int rank)
/*
 EXCHANGE trades edge values with up to four neighbors.
 */
{
	double ex0[nodeedge];
	double ex1[nodeedge];
	double ex2[nodeedge];
	double ex3[nodeedge];
	int i;
	double in0[nodeedge];
	double in1[nodeedge];
	double in2[nodeedge];
	double in3[nodeedge];
	int partner;

	MPI_Request requests[8];
	MPI_Status status[8];
	int tag;

	for (i = 0; i < 8; i++) {
		requests[i] = MPI_REQUEST_NULL;
	}
	/*
	 Receive from UP neighbor (0).
	 */
	if (comm[0] == 1) {
		int partnerCoords[2] = { coords[0] - 1, coords[1] };
		MPI_Cart_rank(cartcomm, partnerCoords, &partner);
		tag = 0;
		MPI_Irecv(&in0, nodeedge, MPI_DOUBLE, partner, tag, MPI_COMM_WORLD,
				&requests[0]);
	}
	/*
	 Receive from RIGHT neighbor (1).
	 */
	if (comm[1] == 1) {
		int partnerCoords[2] = { coords[0], coords[1] + 1 };
		MPI_Cart_rank(cartcomm, partnerCoords, &partner);
		tag = 1;
		MPI_Irecv(&in1, nodeedge, MPI_DOUBLE, partner, tag, MPI_COMM_WORLD,
				&requests[1]);
	}
	/*
	 Receive from DOWN neighbor (2).
	 */
	if (comm[2] == 1) {
		int partnerCoords[2] = { coords[0] + 1, coords[1] };
		MPI_Cart_rank(cartcomm, partnerCoords, &partner);
		tag = 2;
		MPI_Irecv(&in2, nodeedge, MPI_DOUBLE, partner, tag, MPI_COMM_WORLD,
				&requests[2]);
	}
	/*
	 Receive from LEFT neighbor (3).
	 */
	if (comm[3] == 1) {
		int partnerCoords[2] = { coords[0], coords[1] - 1 };
		MPI_Cart_rank(cartcomm, partnerCoords, &partner);
		tag = 3;
		MPI_Irecv(&in3, nodeedge, MPI_DOUBLE, partner, tag, MPI_COMM_WORLD,
				&requests[3]);
	}
	/*
	 Send up from DOWN (2) neighbor.
	 */
	if (comm[0] == 1) {
		int partnerCoords[2] = { coords[0] - 1, coords[1] };
		MPI_Cart_rank(cartcomm, partnerCoords, &partner);
		tag = 2;
		setex(ex0, M, 0);
		MPI_Isend(&ex0, nodeedge, MPI_DOUBLE, partner, tag, MPI_COMM_WORLD,
				&requests[4]);
	}
	/*
	 Send right form LEFT (3) neighbor.
	 */
	if (comm[1] == 1) {
		int partnerCoords[2] = { coords[0], coords[1] + 1 };
		MPI_Cart_rank(cartcomm, partnerCoords, &partner);
		tag = 3;
		setex(ex1, M, 1);
		MPI_Isend(&ex1, nodeedge, MPI_DOUBLE, partner, tag, MPI_COMM_WORLD,
				&requests[5]);
	}
	/*
	 Send down from UP (0) neighbor.
	 */
	if (comm[2] == 1) {
		int partnerCoords[2] = { coords[0] + 1, coords[1] };
		MPI_Cart_rank(cartcomm, partnerCoords, &partner);
		tag = 0;
		setex(ex2, M, 2);
		MPI_Isend(&ex2, nodeedge, MPI_DOUBLE, partner, tag, MPI_COMM_WORLD,
				&requests[6]);
	}
	/*
	 Send left from RIGHT (1) neighbor.
	 */
	if (comm[3] == 1) {
		int partnerCoords[2] = { coords[0], coords[1] - 1 };
		MPI_Cart_rank(cartcomm, partnerCoords, &partner);
		tag = 1;
		setex(ex3, M, 3);
		MPI_Isend(&ex3, nodeedge, MPI_DOUBLE, partner, tag, MPI_COMM_WORLD,
				&requests[7]);
	}

	MPI_Waitall(8, requests, status);
	/*
	 Copy boundary values, sent by neighbors, into M.
	 */
	if (comm[0] == 1) {
		unpack(M, 0, in0);
	}
	if (comm[1] == 1) {
		unpack(M, 1, in1);
	}
	if (comm[2] == 1) {
		unpack(M, 2, in2);
	}
	if (comm[3] == 1) {
		unpack(M, 3, in3);
	}

	return;
}

void addHeat(double M[][nodeedge + 2], int it, int coords[]) {
	int i;
	int j;
	for (i = 1; i <= nodeedge; i++) {
		for (j = 1; j <= nodeedge; j++) {
			M[i][j] += heatFunction(i+coords[0]*nodeedge, j+coords[1]*nodeedge, it)/4.0;
		}
	}
}

void timeIteration(double M[][nodeedge + 2], double result[][nWithBorders], int rank,
		int comm[], int step, int finish, int coords[]) {
	int count;
	int i;
	double in;
	int index;
	int it;
	int j;
	int k;
	int l;
	FILE *f;
	double MM[n * n];
	double send[nodeedge][nodeedge];
	for (it = 1; it <= finish; it++) {
		if (rank == 0){
			char filename[20];
			sprintf(filename,"wynik-%d",it);
			f = fopen(filename, "wt");
		}
		dojacobi(M);

		addHeat(M, it, coords);
		exchange(M, comm, rank);
		if (it % step == 0) {
			for (i = 0; i < nodeedge; i++) {
				for (j = 0; j < nodeedge; j++) {
					send[i][j] = M[i + 1][j + 1];
				}
			}

			count = nodeedge * nodeedge;

			MPI_Gather(&send, count, MPI_DOUBLE, &MM, count, MPI_DOUBLE, 0,
					MPI_COMM_WORLD);

			printf("  ITERATE gathered updated results to process 0.\n");

			if (rank == 0) {
				index = 0;
				for (k = 0; k < nblock; k++) {
					for (l = 0; l < nblock; l++) {
						for (i = k * nodeedge; i < (k + 1) * nodeedge; i++) {
							for (j = l * nodeedge; j < (l + 1) * nodeedge;
									j++) {
								result[i + 1][j + 1] = MM[index];
								index++;
							}
						}
					}
				}
				printf("did %i iterations\n", it);
				for (i = 0; i < nWithBorders; i++) {
					for (j = 0; j < nWithBorders; j++) {
						fprintf(f, "%lf ", result[i][j]);
					}
					fprintf(f, "\n");
				}

				fprintf(f, "\n");
			}

		}

	}
	if (rank == 0)
		fclose(f);
	return;
}

void iterate(double M[][nodeedge + 2], double result[][nWithBorders], int rank,
		int comm[])

/*
 ITERATE controls the iteration, including convergence checking.
 */
{
	int count;
	double diff;
	int done;
	double ediff;
	int i;
	double in;
	int index;
	int it;
	int j;
	int k;
	int l;
	double MM[n * n];
	double mold[nodeedge + 2][nodeedge + 2];
	double send[nodeedge][nodeedge];

	it = 0;
	done = 0;
	for (i = 1; i <= nodeedge; i++) {
		for (j = 1; j <= nodeedge; j++) {
			mold[i][j] = M[i][j];
		}
	}

	while (done == 0) {
		it++;
		/*
		 Exchange values with neighbors, update red squares, exchange values
		 with neighbors, update black squares.
		 */
		if (method == 0) {
			dojacobi(M);
			exchange(M, comm, rank);
		} else {
			exchange(M, comm, rank);
			dored(M);
			exchange(M, comm, rank);
			doblack(M);
		}
		/*
		 Check for convergence every 20 iterations.
		 Find the average absolute change in elements of M.
		 Maximum iterations is 5000.
		 */
		if (5000 < it) {
			done = 1;
		}

		if ((it % 20 == 0) && (done != 1)) {
			diff = 0.0;
			for (i = 1; i <= nodeedge; i++) {
				for (j = 1; j <= nodeedge; j++) {
					ediff = M[i][j] - mold[i][j];
					if (ediff < 0.0) {
						ediff = -ediff;
					}
					diff = diff + ediff;
					mold[i][j] = M[i][j];
				}
			}
			diff = diff / ((double) (nodeedge * nodeedge));
			/*
			 IN = sum of DIFF over all processes.
			 */
			MPI_Allreduce(&diff, &in, 1, MPI_DOUBLE, MPI_SUM, MPI_COMM_WORLD);

			if (in < (double) nproc * 0.00001) {
				done = 1;
			}
		}
	}
	/*
	 Send results to task 0.
	 */
	for (i = 0; i < nodeedge; i++) {
		for (j = 0; j < nodeedge; j++) {
			send[i][j] = M[i + 1][j + 1];
		}
	}

	count = nodeedge * nodeedge;

	MPI_Gather(&send, count, MPI_DOUBLE, &MM, count, MPI_DOUBLE, 0,
			MPI_COMM_WORLD);

	printf("  ITERATE gathered updated results to process 0.\n");
	/*
	 Storage on task 0 has to be consistent with a NBLOCK x NBLOCK decomposition.

	 I believe the array form of RESULT is only needed at the end of the
	 program (and not even then, really).  So we could probably skip this
	 work of rearranging the data here.  JVB, 11 January 2012.
	 */
	if (rank == 0) {
		printf("did %i iterations\n", it);

		index = 0;

		for (k = 0; k < nblock; k++) {
			for (l = 0; l < nblock; l++) {
				for (i = k * nodeedge; i < (k + 1) * nodeedge; i++) {
					for (j = l * nodeedge; j < (l + 1) * nodeedge; j++) {
						result[i + 1][j + 1] = MM[index];
						index++;
					}
				}
			}
		}
	}
	return;
}

void setcomm(int coords[], int dims[], int comm[])

/*
 SETCOMM determines the active communication directions.
 */
{
	int i;

	for (i = 0; i < 4; i++) {
		comm[i] = 1;
	}
	/*
	 Up neighbor?
	 */
	if (coords[0] == 0) {
		comm[0] = 0;
	}
	/*
	 Right neighbor?
	 */
	if (coords[1] == dims[1] - 1) {
		comm[1] = 0;
	}
	/*
	 Down neighbor?
	 */
	if (coords[0] == dims[0] - 1) {
		comm[2] = 0;
	}
	/*
	 Left neighbor?
	 */
	if (coords[1] == 0) {
		comm[3] = 0;
	}

	return;
}

void setex(double ex[], double M[][nodeedge + 2], int which)
/*
 SETEX pulls off the edge values of M to send to another task.
 */
{
	int i;

	switch (which) {
	case 0: {
		for (i = 1; i <= nodeedge; i++) {
			ex[i - 1] = M[1][i];
		}
		break;
	}
	case 1: {
		for (i = 1; i <= nodeedge; i++) {
			ex[i - 1] = M[i][nodeedge];
		}
		break;
	}
	case 2: {
		for (i = 1; i <= nodeedge; i++) {
			ex[i - 1] = M[nodeedge][i];
		}
		break;
	}
	case 3: {
		for (i = 1; i <= nodeedge; i++) {
			ex[i - 1] = M[i][1];
		}
		break;
	}
	}
	return;
}

void unpack(double M[][nodeedge + 2], int where, double in[])

/*
 UNPACK puts the vector of new edge values into the edges of M.
 */
{
	int i;

	if (where == 0) {
		for (i = 0; i < nodeedge; i++) {
			M[0][i + 1] = in[i];
		}
	} else if (where == 1) {
		for (i = 0; i < nodeedge; i++) {
			M[i + 1][nodeedge + 1] = in[i];
		}
	} else if (where == 2) {
		for (i = 0; i < nodeedge; i++) {
			M[nodeedge + 1][i + 1] = in[i];
		}
	} else if (where == 3) {
		for (i = 0; i < nodeedge; i++) {
			M[i + 1][0] = in[i];
		}
	}

	return;
}
